%% Load the data - Select the Files: 'ViscotaxisExperiments', 'ViscosityProfiles', 'SimulationDensitySpeedProfiles'
close all; clear all; clc; 
clearvars;
total_directory=uigetdir(); %pick the folder you want to work in 
total_directory = dir(total_directory);
dirFlags = [total_directory.isdir];
total_directory = total_directory(~ismember({total_directory.name},{'.','..'}));
cd([total_directory(1).folder])
%% Combine Tracks
tracks_ExpComb = []; c = 0;h = waitbar(0,'Calculating ... ');
for i = 1:length(total_directory)
    clear tracks
    load([total_directory(i).folder filesep total_directory(i).name]);
    tracks_ExpComb{i} = tracks;
end
waitbar(c+1 / 5,h);
% Filter Tracks
TracksSel = [];YwallMax_s = [];YwallMin_s = [];myparameters = [];
myparameters.lengthdev = 4;

for i = 1:length(tracks_ExpComb)
    clear InputTracks tracks_c tracks_filt
    InputTracks = tracks_ExpComb{1,i};
    [tracks_c] = Camera2MicronSecond(InputTracks,mag,fps,Camera_um_pix);
    [tracks_filt, trackstats] = FilterTracks(tracks_c, myparameters);
    YwallMin_s(1,i) = min(vertcat(tracks_filt(:).y_um)); YwallMax_s(1,i) = max(vertcat(tracks_filt(:).y_um));
    TracksSel{i} = tracks_filt;
end
waitbar(c+1 / 5,h);
Nbins = 11;
YwallMin = min(YwallMin_s); YwallMax = max(YwallMax_s);
%
results_Exp = [];
for i = 1:length(TracksSel)
    clear results InputTracks
    InputTracks = TracksSel{1,i};
    [results, ListYs, BinSize] = SpeedandPDF_Final(InputTracks, fps, mag, Nbins,YwallMin,YwallMax);
    results_Exp{i} = results
end
waitbar(c+1 / 5,h);
%
if strcmp(PerPEO,'0.5');RevExp = 2; else; RevExp = 0; end
[results,FinalResults] = ReplicateExpCombine(results_Exp,RevExp);
RevExp
waitbar(c+1 / 5,h); close(h)
'Done'
%
Vel_rho = FinalResults.Speed.*FinalResults.Density;
Vel_rho_STDErrorPro = sqrt((FinalResults.Speed.*FinalResults.Density_STDError).^2 + (FinalResults.Density.*FinalResults.Speed_STDError).^2);

start = 4;fin = 10;
PerPEOColor = [0 0 0;               % Control 
                0.635 0.078 0.184;  % 0.05 % PEO
                0.588 0.286 0.612;  % 0.1% PEO
                0.929 0.694 0.125;  % 0.25% PEO
                0.466 0.674 0.188;  % 0.5% PEO
                0.85 0.325 0.0980]; % 0.82% PEO
if strcmp(PerPEO,'0.1')
    load('.\ViscosityProfiles\0pt1PEO')
    vis = cell2mat(mu'); LineColor = PerPEOColor(3,:); SimIndex = 2;
elseif strcmp(PerPEO,'0.25')
    load('.\ViscosityProfiles\0pt25PEO')
    vis = cell2mat(mu'); LineColor = PerPEOColor(4,:); SimIndex = 3;
elseif strcmp(PerPEO,'0.5')
    load('.\ViscosityProfiles\0pt5PEO')
    vis = cell2mat(mu'); LineColor = PerPEOColor(5,:); SimIndex = 4;
elseif strcmp(PerPEO,'0.82')
    load('.\ViscosityProfiles\0pt82PEO')
    vis = cell2mat(mu'); LineColor = PerPEOColor(6,:); SimIndex = 5;
else
    load('.\ViscosityProfiles\0pt0PEO')
    vis = cell2mat(mu'); start = 1; fin = 5; PerPEOColor = [1 1 1]; LineColor = [0 0 0]; SimIndex = 1;
end

x = FinalResults.Pos; x = x - (median(x)-500);
y = Vel_rho;error = Vel_rho_STDErrorPro;
xPos_Visc = xPos + (500 - median(xPos)); MeanVis = mean(vis(start:fin,:));

ExpVisc_inter = interp1(xPos_Visc, MeanVis, linspace(x(1),x(end),length(x)), 'linear', 'extrap');
m = -0.934; k = 4.635; % Value from fitting power function to bulk swimming speed exp.
FluxMag = exp(k)/trapz(x,ExpVisc_inter.^(-m)); %Flux Mag = V_0*rho_0 = rho(x)*V(x)

%Load Simulation Data
load('.\SimulationDensitySpeedProfiles\SimulationDensityandVelocity.mat');
SimDensity = cell2mat(sim_den'); SimSpeed = cell2mat(sim_v'); SimPos = cell2mat(sim_x');

if SimIndex > 0
    x_sim = SimPos(SimIndex,:); rho_sim = SimDensity(SimIndex,:);
    V_sim = SimSpeed(SimIndex,:); y_Flux_sim = rho_sim.*V_sim;
end

Marker = '-';
FillColor = LineColor;
LineWidth = 1;
alpha_errorbar = 0.50;
alw = 0.5; % Border Width
SchnitzerDeparture_Exp = figure;
    plot(x,y,'-','color',LineColor,'LineWidth',LineWidth)
    hold on;
    patch = fill([x,fliplr(x)],[y + error,fliplr(y - error)],FillColor);
    set(patch, 'edgecolor', 'none'); set(patch, 'FaceAlpha', alpha_errorbar);
    p1 = line([0 1000],FluxMag*ones(2,1)); p1.Color = [0.6 0.6 0.6]; p1.LineStyle = '-'; p1.LineWidth = 0.5*LineWidth
    if SimIndex > 0
        plot(x_sim,y_Flux_sim,':','Color',LineColor,'LineWidth',LineWidth*1.5)
    end
    uistack(p1,'bottom')
    
    hold off; 
    xlim([0 1000]);ylim([0 .18])
    box on;set(gca, 'Layer', 'top');
    xticks([0,250,500,750,1000]);
    yticks([0 0.05 0.1 0.15]);
      
    xticklabels({'','','',''});yticklabels({'','','',''});ylabel('');xlabel('');title('')
    mypapersize = [.9 .9]; %[width height]
    set(gca,'LineWidth',alw);
    set(gca,'XColor',[0 0 0]);set(gca,'YColor',[0 0 0]);
    set(gcf, 'PaperUnits', 'inches', 'papersize', mypapersize, 'PaperPosition', [0 0 mypapersize]);
    set(SchnitzerDeparture_Exp,'Units','inches','Position',[0 0 mypapersize]);
    ax = gca;outerpos = ax.OuterPosition;ti = ax.TightInset;
    left = outerpos(1);bottom = outerpos(2);ax_width = outerpos(3);ax_height = outerpos(4);
    ax.Position = [left+0.011 bottom+0.01 ax_width-0.025 ax_height-0.022];
% Save Figure
    figdir = '.\Figure1\Schnitzer';
    print(SchnitzerDeparture_Exp, '-painters','-dpdf', '-r1200', [figdir,'\Rho_V',PerPEO,'.pdf']);
% Save Data in Excel
dir_Excel = '.\Source data\MainText'
Tab = 10;
if strcmp(PerPEO,'0.1'); column1 = 'I4'; column2 = 'J4'; column3 = 'K4'; column4 = 'L4'; column5 = 'M4'; column6 = 'N4';column7 = 'O4';
    elseif strcmp(PerPEO,'0.25'); column1 = 'Q4'; column2 = 'R4'; column3 = 'S4'; column4 = 'T4'; column5 = 'U4'; column6 = 'V4';column7 = 'W4';
    elseif strcmp(PerPEO,'0.5'); column1 = 'Y4'; column2 = 'Z4'; column3 = 'AA4'; column4 = 'AB4'; column5 = 'AC4'; column6 = 'AD4';column7 = 'AE4';
    elseif strcmp(PerPEO,'0.82'); column1 = 'AG4'; column2 = 'AH4'; column3 = 'AI4'; column4 = 'AJ4'; column5 = 'AK4'; column6 = 'AL4';column7 = 'AM4';
    else; column1 = 'A4'; column2 = 'B4'; column3 = 'C4'; column4 = 'D4'; column5 = 'E4'; column6 = 'F4'; column7 = 'G4';
end
writematrix(x',[dir_Excel,'Figure1.xlsx'],'Sheet',Tab,'Range',column1)
writematrix(y',[dir_Excel,'Figure1.xlsx'],'Sheet',Tab,'Range',column2)
writematrix(error',[dir_Excel,'Figure1.xlsx'],'Sheet',Tab,'Range',column3)
writematrix(linspace(0,1000,11)',[dir_Excel,'Figure1.xlsx'],'Sheet',Tab,'Range',column4)
writematrix(FluxMag*ones(11,1),[dir_Excel,'Figure1.xlsx'],'Sheet',Tab,'Range',column5)
 
writematrix(x_sim',[dir_Excel,'Figure1.xlsx'],'Sheet',Tab,'Range',column6)
writematrix(y_Flux_sim',[dir_Excel,'Figure1.xlsx'],'Sheet',Tab,'Range',column7)   
%% FUNCTION: Convert pix/frame to um/s
function [c_tracks] = Camera2MicronSecond(tracks,mag,fps,Camera_micron_pix)
% Camera2MicronSecond: Converts camera units (frames and pix) to physical
% units (microns and seconds). Cameras have different pixel dimensions. 
% For example, the Zyla has 6.5 microns per pixel
    c_tracks = tracks;
    for i=1:length(c_tracks)
        c_tracks(i).t_sec = c_tracks(i).t/fps;
        c_tracks(i).x_um = c_tracks(i).x.*Camera_micron_pix/mag; c_tracks(i).y_um = c_tracks(i).y.*Camera_micron_pix/mag;
        c_tracks(i).u_um = c_tracks(i).u.*Camera_micron_pix.*fps./mag; c_tracks(i).v_um = c_tracks(i).v.*Camera_micron_pix.*fps./mag;    
        % Calculate the Speed of the Trajectory (magnitude of u-velocity and v-velocity)    
        s = sqrt(c_tracks(i).v.^2+c_tracks(i).u.^2); c_tracks(i).Speed = s;
        s_um = sqrt(c_tracks(i).v_um.^2+c_tracks(i).u_um.^2);  c_tracks(i).Speed_um = s_um;   
    end
end
%% FUNCTION: Filter Tracks
function [f_tracks, trackstats] = FilterTracks(tracks, myparameters)

if ~isfield(myparameters,'minmaxlengthtime')||isempty(myparameters.minmaxlengthtime); myoptions.minmaxlengthtime=0; else; minlengthtime = myparameters.minmaxlengthtime(1);maxlengthtime=myparameters.minmaxlengthtime(2); myoptions.minmaxlengthtime=1; end
if ~isfield(myparameters,'minmedianspeed')||isempty(myparameters.minmedianspeed);  myoptions.minmedianspeed=0; else; minmedianspeed = myparameters.minmedianspeed; myoptions.minmedianspeed=1; end
if ~isfield(myparameters,'minRG') || isempty(myparameters.minRG); myoptions.minRG = 0; else; minRG = myparameters.minRG; myoptions.minRG = 1; end
if ~isfield(myparameters,'lengthdev')||isempty(myparameters.lengthdev); myoptions.lengthdev=0; else; lengthdev=myparameters.lengthdev; myoptions.lengthdev=1; end
if ~isfield(myparameters,'minmaxcontourlength')||isempty(myparameters.minmaxcontourlength);myoptions.minmaxcontourlength=0; else; mincontourlength=myparameters.minmaxcontourlength(1); maxcontourlength=myparameters.minmaxcontourlength(2); myoptions.minmaxcontourlength=1; end

f_tracks=tracks;

if myoptions.minmaxlengthtime==1
    ind_minmaxlengthtime=zeros(1,length(f_tracks));tracklength=zeros(1,length(f_tracks));
    for k=1:length(f_tracks)
        tracklength(k) = length(f_tracks(k).x);
        if tracklength(k) < maxlengthtime && tracklength(k) > minlengthtime
            ind_minmaxlengthtime(1,k)=1;
        end
    end
    trackstats.tracklengthtime=tracklength;
else
    ind_minmaxlengthtime=ones(1,length(f_tracks));
end

if myoptions.minmedianspeed==1
    med=zeros(1,length(f_tracks));ind_minmedianspeed=zeros(1,length(f_tracks));
    for k=1:length(f_tracks)
        u_r = sqrt(tracks(k).u.^2 + tracks(k).v.^2);
        med(k) = median(u_r);
        if med(k) > minmedianspeed
            ind_minmedianspeed(1,k)=1;
        end
    end
    trackstats.medianspeed=med;
else
    ind_minmedianspeed=ones(1,length(f_tracks));
end

if myoptions.minRG==1
    ind_minRG=zeros(1,length(f_tracks));RG=zeros(1,length(f_tracks));
    for k=1:length(f_tracks)
        RG(1,k)= sqrt((max(f_tracks(k).x)-min(f_tracks(k).x))^2+(max(f_tracks(k).y)-min(f_tracks(k).y))^2);
        if RG(1,k)>minRG
            ind_minRG(1,k)=1;
        end
    end
    trackstats.RG=RG;
else
    ind_minRG=ones(1,length(f_tracks));
    
end

if myoptions.lengthdev==1
    ind_lengthdev=zeros(1,length(f_tracks));RG=zeros(1,length(f_tracks));contourlength=zeros(1,length(f_tracks));dev=zeros(1,length(f_tracks));
    s_tracks=EqualSpace(f_tracks,1);
    for k=1:length(f_tracks)
        RG(1,k)= sqrt((max(f_tracks(k).x)-min(f_tracks(k).x))^2+(max(f_tracks(k).y)-min(f_tracks(k).y))^2);
        contourlength(1,k)=length(s_tracks(k).x);
        dev(1,k)=contourlength(1,k)/RG(1,k);
        if dev(1,k)<lengthdev
            ind_lengthdev(1,k)=1;
        end
    end
    trackstats.RG=RG;
    trackstats.contourlength=contourlength;
    trackstats.dev=dev;
else
    ind_lengthdev=ones(1,length(f_tracks));
    
end

if myoptions.minmaxcontourlength==1
    ind_minmaxcontourlength=zeros(1,length(f_tracks));
    if myoptions.lengthdev==0
        s_tracks=EqualSpace(f_tracks,1);contourlength=zeros(1,length(f_tracks));
        for k=1:length(s_tracks)
            contourlength(1,k)=length(s_tracks(k).x);
            if contourlength(1,k)> mincontourlength && contourlength(1,k)<maxcontourlength
                ind_minmaxcontourlength(1,k)=1;
            end
        end
        trackstats.contourlength=contourlength;
    else
        for k=1:length(contourlength)
            if contourlength(1,k)> mincontourlength && contourlength(1,k)< maxcontourlength
                ind_minmaxcontourlength(1,k)=1;
            end
        end
    end 
else
    ind_minmaxcontourlength=ones(1,length(f_tracks));
end

ind_total=ind_minmaxlengthtime & ind_minmedianspeed & ind_minRG & ind_lengthdev & ind_minmaxcontourlength;

f_tracks=f_tracks(ind_total);
end

function tracks_lambda=EqualSpace(tracks,spacing)
tracks_lambda=struct('x',[],'y',[],'s',[],'u',[],'v',[],'t',[]);
for i=1:length(tracks)
    s=zeros(1,length(tracks(i).x));
    for ii=2:length(tracks(i).x)
        s(ii)=s(ii-1)+sqrt((tracks(i).x(ii)-tracks(i).x(ii-1))^2+(tracks(i).y(ii)-tracks(i).y(ii-1))^2);
    end
    sq=0:spacing:s(end);
    x=interp1(s,tracks(i).x,sq);y=interp1(s,tracks(i).y,sq);u=interp1(s,tracks(i).u,sq,'pchip');v=interp1(s,tracks(i).v,sq,'pchip');t=interp1(s,tracks(i).t,sq);
    tracks_lambda(i,1).x(:,1)=x;tracks_lambda(i,1).y(:,1)=y;tracks_lambda(i,1).u(:,1)=u;tracks_lambda(i,1).v(:,1)=v;tracks_lambda(i,1).s(:,1)=s;tracks_lambda(i,1).t(:,1)=t;
end
end
%% FUNCTION: Bin the tracks to get Spatial Density and Swimming Speed
function [results, ListYs, BinSize] = SpeedandPDF_Final(tracks_in, fps, mag, Nbins,YwallMin,YwallMax)
%SpeedandPDF_Final: Using the tracks from the Guasto Lab tracking
%code (tracks) we can bin the y-position with corresponding speeds to
%create a PDF and speed distribution of the tracks in the channel.

    tracks = tracks_in;
    Ys= cat(1,tracks(:).y_um);  Ys_um = Ys;   
    Speed = cat(1,tracks(:).Speed_um);
    Urs = cat(1,tracks(:).Speed_um); U_um = Urs;
    
    %Determine the minumum and maximum y values
    Ymin = YwallMin; Ymax = YwallMax;

    ListYs = linspace(Ymin,Ymax,Nbins+1);   %Create the y-intervals to categorize the speeds 
    Counts  = zeros(1,Nbins);               %create the matrix to store the number of data in the section
    Speeds  = Counts;                       %Initialize the variable speeds and counts
    Speeds_um  = Counts;
    SpeedsSquare = Counts;                  %Initialize the variable SpeedsSquare and Counts
    SpeedsSquare_um = Counts;
    BinSize = ListYs(2)-ListYs(1);          %Determine the size of each bin

    for jj = 1:length(Ys)                                              % loop to run through all the y values from Ys
        Yjj = Ys(jj);                                                  % Assigns the variable Yjj to the actual y value
        Threshold = (Yjj-ListYs(1))/BinSize;                           % Determine which bin the y position falls into
        RowY = ceil(Threshold);                                        % Rounds the Threshold up to assign it to the bin number
        if (0<RowY)&&(RowY<=Nbins)                                     % Loop to check if the Ys value falls into selected bin
        Counts(RowY) = Counts(RowY)+1;                                 % Counter to track the number of data points in each bin
        Speeds(RowY) = Speeds(RowY)+Urs(jj);                           % Counter for the average calculation below (sum of the speeds for each bin)
        SpeedsSquare(RowY) = SpeedsSquare(RowY) + Urs(jj)^2;           % Counter for the Variance equation
        
        Speeds_um(RowY) = Speeds_um(RowY)+U_um(jj);                    % um/sec: Counter for the average calculation below (sum of the speeds for each bin)
        SpeedsSquare_um(RowY) = SpeedsSquare_um(RowY) + U_um(jj)^2;    % um/sec: Counter for the Variance equation
        
        end
    end

    AverageSpeeds = Speeds./Counts;     %Calculate the average speed for each bin
    AverageSpeeds_um = Speeds_um./Counts;
    
    VarianceSpeeds = SpeedsSquare./Counts-AverageSpeeds.^2;     %Calculate the variance
    VarianceSpeeds_um = SpeedsSquare_um./Counts-AverageSpeeds_um.^2;
    d = 1;
    results(d).Speeds = AverageSpeeds; results(d).Errors = VarianceSpeeds;
    results(d).sqrVariance = sqrt(VarianceSpeeds); results(d).Ybins = ListYs(1:end-1)+BinSize/2;
    results(d).InverseMeanSpeed = 1./AverageSpeeds; results(d).ErrorBar = sqrt(VarianceSpeeds)/Counts;
    results(d).Counts = Counts;
    
    [pdf,ybins] = hist(Ys,Nbins); trapz(ybins,pdf); pdf = pdf/trapz(ybins,pdf);
    results(d).PDF = pdf; results(d).YbinsPDF = ybins;
    
    [pdf_speed,ybins_s] = hist(Speed,Nbins); trapz(ybins_s,pdf_speed); pdf_speed = pdf_speed/trapz(ybins_s,pdf_speed);
    results(d).PDFSpeed = pdf_speed; results(d).YbinsSpeed = ybins_s;
    
    results(d).Speeds_um = AverageSpeeds_um; results(d).Errors_um = VarianceSpeeds_um; results(d).sqrVariance_um = sqrt(VarianceSpeeds_um);
    results(d).InverseMeanSpeed_um = 1./AverageSpeeds_um; results(d).ErrorBar_um = sqrt(VarianceSpeeds_um)/Counts;
    
    [pdf_um,ybins_um] = hist(Ys_um,Nbins); trapz(ybins_um,pdf_um);
    pdf_um = pdf_um/trapz(ybins_um,pdf_um);
    results(d).PDF_um = pdf_um; results(d).YbinsPDF_um = ybins_um;

end

%% FUNCTION: Combine replicate experiments
function [results,FinalResults] = ReplicateExpCombine(results_Exp_in,RevExp)
% ReplicateExpCombine: Combine the replicate experiments for each viscosity profile and
% calculate the mean spatial swimming speed and PDF. Return the values
% for making figures. For 5 cP viscotaxis experiment, Exp2, flip the data
% due to the gradient being the other way.

results_Exp = results_Exp_in;
results = [];FinalResults = [];
for i = 1:length(results_Exp)
    results(i).Pos = results_Exp{1,i}.Ybins;
    if i == RevExp
        results(i).Density = flip(results_Exp{1,i}.PDF_um); results(i).Speed = flip(results_Exp{1,i}.Speeds_um);
    else
        results(i).Density = results_Exp{1,i}.PDF_um; results(i).Speed = results_Exp{1,i}.Speeds_um;
    end
end
FinalResults.Pos = mean(vertcat(results.Pos));
Store_MeanExpPDF = vertcat(results.Density); FinalResults.Density = mean(Store_MeanExpPDF);
FinalResults.Density_STD = nanstd(Store_MeanExpPDF); FinalResults.Density_STDError = nanstd(Store_MeanExpPDF)./sqrt(length(results_Exp));
Store_MeanExpSpeed = vertcat(results.Speed); FinalResults.Speed = mean(Store_MeanExpSpeed);
FinalResults.Speed_STD = nanstd(Store_MeanExpSpeed); FinalResults.Speed_STDError = nanstd(Store_MeanExpSpeed)./sqrt(length(results_Exp));
end